"""5.5 【资源类】Resource 客户端开发"""
import asyncio
import json
from contextlib import AsyncExitStack
from typing import Optional

from mcp import ClientSession, ListResourcesResult
from openai import OpenAI
from mcp.client.sse import sse_client

class MCPClient:

    def __init__(self):
        #创建线程管理栈
        self.exit_stack = AsyncExitStack()
        self.session: Optional[ClientSession] = None
        self.deepseek = OpenAI(
            api_key="sk-15af4e21f828460683b16ce9e78b2346",
            base_url="https://api.deepseek.com"
        )
        self.resources = {}

    """创建连接服务端"""
    async def connect_to_server(self, server_path: str):
        # 一、创建StdioServerParameter参数信息
        # server_parameters = StdioServerParameters(
        #     command="python",
        #     args=[server_path],
        #     env=None
        # )
        # 二、 创建sse_client
        client = sse_client(url=server_path)
        stdio_transport = await self.exit_stack.enter_async_context(client)
        read_stream, write_stream = stdio_transport

        # 三、创建ClientSession
        client_session = ClientSession(read_stream, write_stream)
        self.session = await self.exit_stack.enter_async_context(client_session)
        # 四、初始化session
        await self.session.initialize()


    async def execute(self,query:str):

        # 一、通过session列表所有的工具 及组装function_calling
        ''' 获取资源类的信息 '''
        list_resource_result:ListResourcesResult = await self.session.list_resources()
        resources = list_resource_result.resources
        print("\nConnected to server with resource:",resources)
        resource_tools = []
        for resource in resources:
            uri = resource.uri
            name = resource.name
            description = resource.description
            mime_type = resource.mimeType
            input_schema = None
            print(f"URI: {uri}, Name: {name}, Description: {description}, Input Schema: {input_schema}")
            # 保存资源类信息，为后续function_calling做准备
            self.resources[name] = {
                "uri":uri,
                "name":name,
                "description":description,
                "mime_type":mime_type,
                "input_schema":input_schema,
            }
            # 保存资源类信息，为后续function_calling做准备
            resource_tools.append({
                "type":"function",
                "function":{
                    "name":name,
                    "description":description,
                    "parameters":input_schema
                }
            })


        # 组装function calling

        # 二、 大模型调用参数：组装messages
        messages = [
            {
                "role": "user",
                "content": query
            }
        ]
        # 第一次调用大模型做决策（选择合适的资源类）
        deepseek_response = self.deepseek.chat.completions.create(
            model="deepseek-chat",
            messages=messages,
            tools=resource_tools,
        )
        # 获取大模型的决策结果
        print("==== deepseek 决策结果：",deepseek_response)
        choice_result = deepseek_response.choices[0]
        print("最终选择的工具：",choice_result.message.tool_calls[0].function.name)

        '''  ******************* function calling 过程 *******************'''
        if choice_result.finish_reason == "tool_calls":
            # 根据大模型选择的工具，进行调用执行
            messages.append(choice_result.message.model_dump())
            # 调用工具链
            tool_call = choice_result.message.tool_calls[0]
            function_name = tool_call.function.name
            arguments = json.loads(tool_call.function.arguments)


            ''' ************** 调用资源类 *************'''
            uri = self.resources[function_name]['uri']
            #tool_result = await self.session.call_tool(name = function_name,arguments=arguments)
            resource_result = await self.session.read_resource(uri)

            print("==== 工具调用结果元数据：",resource_result)
            # 把 resource_result.content改为resource_result.contents[0].text
            print("====  工具计算结果：",resource_result.contents[0].text)

            ''' ************** 3.5 最终推理结果 (第二次调用大模型）*************'''
            # 组装第二次调用大模型的message参数:角色需要为tool,来响应工具的调用
            messages.append({
                "role":"tool",
                "content":resource_result.contents[0].text,
                "tool_call_id":tool_call.id
            })
            #再次调用大模型
            deepseek_response = self.deepseek.chat.completions.create(
                model="deepseek-chat",
                messages=messages,
            )
            # 获取最终的结果
            result = deepseek_response.choices[0].message.content
            print("==== 最终的结果：",result)


    #关闭连接
    async def cleanup(self):
        await self.exit_stack.aclose()

async def main():
    client = MCPClient()
    while True:
        #输入文本内容
        query = input("请输入咨询内容：")
        if query == "exit":
            break
        try:
            await client.connect_to_server("http://localhost:8000/sse")
            #如果没有输入,query,则默认查询id为1001的书籍信息
            if not query:
                query = "查询id为1001的书籍信息"
            await client.execute(query)
        except Exception as e:
            print(f"Error: {str(e)}")
        finally:
            await client.cleanup()

if __name__ == "__main__":
    asyncio.run(main())